#include <stdio.h>
#include <stdlib.h>
#include <iostream>
#include <crs.h>
#include <ell.h>
#include <mtrnd.h>

#ifndef   _HAMILTONIAN_H
#define   _HAMILTONIAN_H

#define W    1
#define CUT_OFF 0.001
#define INDEX(x,y,z,XX,YY,ZZ) ((((z)>((ZZ)-1))?0:(((z)<0)?(ZZ-1):(z)) * (XX) * (YY)) + \
                              (((y)>(YY-1))?0:(((y)<0)?(YY-1):(y)) * (XX)) +  \
                              (((x)>(XX-1))?0:(((x)<0)?(XX-1):(x))))

/*
 #define INDEX_BCC(x,y,z,XX,YY,ZZ) (((z>(ZZ*2-1)?z=0:((z<0)?(z=(ZZ*2-1)):z))/2*(XX)*(YY) + \
								  (y>(YY*2-1)?y=0:((y<0)?(y=(YY*2-1)):y))/2*XX + \
								  (x>(XX*2-1)?x=0:((x<0)?(x=(XX*2-1)):x))/2) + ((z%2==1)?XX*YY*ZZ:0))

 *
 *
#define INDEX_BCC(x,y,z,XX,YY,ZZ) (((z>(ZZ*2-1)?z=0:((z<0)?(ZZ*2-1):z))/2*(XX)*(YY) + \
								  (y>(YY*2-1)?y=0:((y<0)?(YY*2-1):y))/2*XX + \
								  (x>(XX*2-1)?x=0:((x<0)?(XX*2-1):x))/2) + \
								  (((z>(ZZ*2-1)?z=0:((z<0)?(ZZ*2-1):z))%2==1)?XX*YY*ZZ:0))

*/
#define INDEX_BCC(x,y,z,XX,YY,ZZ) (((z>(ZZ*2-1)?0:((z<0)?(ZZ*2-1):z))/2*(XX)*(YY) + \
								  (y>(YY*2-1)?0:((y<0)?(YY*2-1):y))/2*(XX) + \
								  (x>(XX*2-1)?0:((x<0)?(XX*2-1):x))/2) + \
								  (((z>(ZZ*2-1)?0:((z<0)?(ZZ*2-1):z))%2==1)?(XX*YY*ZZ):0))

#define ABS(x) (x>=0?x:(-x))

using namespace std;
template<class T>
struct HamiltonianFile {
  char signature[2]; //"HM"
  unsigned short ver_major;
  unsigned short ver_minor;
  unsigned short ver_build;
  unsigned short precision;   //float is 4, double is 8
  unsigned short entry_type;  //0 is real number, 1 is complex number
  char description[1000];

  int rowsNum;
  int colsNum;
  int elementsNum;
  int* rows;
  int* cols;
  T* elements;
};

template<class T>
class HamiltonianInfo{
private:
	T disorder;
	T a;
	T b;
	int XX;
	int YY;
	int ZZ;

public:
	HamiltonianInfo(){}
	void InitialConfig(int xx, int yy, int zz, T _disorder){
		disorder = _disorder;
		XX = xx; YY = yy; ZZ = zz;
	}
	int getX(){ return XX; }
	int getY(){ return YY; }
	int getZ(){ return ZZ; }
	T getDisorder(){return disorder;}

	void setScale(T _a, T _b)  { a = _a; b = _b; }
	void getScale(T &_a, T &_b){ _a = a; _b = b; }
};

struct ScaleInfo{
	double a;
	double b;
};



class Hamiltonian{

public:
  Hamiltonian();
  ~Hamiltonian();
public:

  template<class T>
  static void normalize(CRS<T>* H, HamiltonianInfo<T> &hi);

  template <class T>
  static CRS<T>* createCubeSingleElectron(HamiltonianInfo<T> &hi, MT::MersenneTwist& mt);

  template <class T>
  static ELL<T>* createCubeSingleElectron_ELL(HamiltonianInfo<T> &hi, MT::MersenneTwist& mt);

  template <class T>
  static CRS<T>* createBCCSingleElectron_CRS(HamiltonianInfo<T> &hi, MT::MersenneTwist& mt);

  template <class T>
  static CRS<T>* readFromFile(char* fileName);

  template <class T>
  static ELL<T>* readFromFile_ELL(char* fileName);

  template <class T>
  void writeCRSToFile(CRS<T>* ha, char* prefix);

  template <class T>
  static ELL<T>* covertToColumnStorage(ELL<T>* c);

  template <class T>
  static ELL<T>* convertToELLFromCRS(CRS<T>* c);

};

using namespace MT;
void OrderData(int a[],int len)
{
  int i,j;
  for(i=len;i>0;i--){
      for(j=0;j<i-1;j++){
          if(a[j] > a[j+1]){
              int temp = a[j];
              a[j] = a[j+1];
              a[j+1] = temp;
          }
      }
  }
}

//normalize matrix using Gershgorin circle theorem
template<typename T>
void Hamiltonian::normalize(CRS<T>* H, HamiltonianInfo<T> &hi){
  int i,j;
  T max=-10,min=10;
  T c,r;
  for(i=0;i<H->rowsNum;i++){
      c = r = 0.0;
      for(j=H->rows[i];j<H->rows[i+1];j++){
          if(H->cols[j] == i){
              c = H->elements[j];
          }else{
              r += ABS(H->elements[j]);
          }
      }
      if(max < c + r){
          max = c + r;
      }
      if(min > c-r){
          min = c - r;
      }
  }

  T a = (max - min) / (2 - CUT_OFF);
  T b = (max + min) / 2;

  for(i=0;i<H->rowsNum;i++){
      for(j=H->rows[i];j<H->rows[i+1];j++){
          if(H->cols[j] == i){
              H->elements[j] -= b;
          }
          H->elements[j] /= a;
      }
  }
  hi.setScale(a, b);
}

template<class T>
T non_on_site_potential(MersenneTwist& mt)
{
  return -1;
  //return (T)rand()*2/RAND_MAX - 1;
}

template<class T>
T on_site_potential(MersenneTwist& mt,T disorder)
{
  return ((T)mt.genrand_real2() - 0.5) * disorder;
}

template<class T>
ELL<T>* Hamiltonian::convertToELLFromCRS(CRS<T>* c){

  int maxCols = 0;
  //find the max number of elements in each row
  for(int i=0;i<c->rowsNum;i++){
      int temp = c->rows[i+1] - c->rows[i];
      if(temp > maxCols)
        maxCols = temp;
  }
  printf("rowsNum:%d\n",c->rowsNum);
  cout<<"Max_Cols:"<<maxCols<<endl;
  printf("rowsNum:%d\n",c->rowsNum);
  printf("Max_Cols:%d\n",maxCols);


  T* elements = (T*)malloc(sizeof(T) * maxCols * c->rowsNum);
  memset(elements, 0, sizeof(T) * maxCols * c->rowsNum);

  int* cols = (int*)malloc(sizeof(int) * maxCols * c->rowsNum);
  memset(cols, 0, sizeof(int) * maxCols * c->rowsNum);

  //copy the data from CRS to ELL
  for(int i=0;i<c->rowsNum;i++){
      memcpy(&elements[i * maxCols],&c->elements[c->rows[i]], sizeof(T) * (c->rows[i+1] - c->rows[i]));
      memcpy(&cols[i * maxCols],&c->cols[c->rows[i]], sizeof(int) * (c->rows[i+1]- c->rows[i]));
  }


  ELL<T>* t = (ELL<T>*)malloc(sizeof(ELL<T>));
  t->colsNum = maxCols;
  t->rowsNum = c->rowsNum;
  t->cols = cols;
  t->elements = elements;

  free(c->elements);
  free(c->cols);
  free(c->rows);
  free(c);

  return t;
}

template <class T>
ELL<T>* Hamiltonian::covertToColumnStorage(ELL<T>* c){
	ELL<T>* t = (ELL<T>*)malloc(sizeof(ELL<T>));
	t->colsNum = c->colsNum;
	t->rowsNum = c->rowsNum;
	t->cols = (int*)malloc(sizeof(int) * t->colsNum * t->rowsNum);
	t->elements = (T*)malloc(sizeof(T) * t->colsNum * t->rowsNum);

	for(int i=0;i<c->colsNum;i++){
		for(int j=0;j<c->rowsNum;j++){
			t->elements[i * c->rowsNum + j] = c->elements[j * c->colsNum + i];
			t->cols[i * c->rowsNum + j] = c->cols[j * c->colsNum + i];
		}
	}
	free(c->cols);
	free(c->elements);
	free(c);
	return t;
}

template<class T>
CRS<T>* Hamiltonian::readFromFile(char* fileName)
{
    FILE* file = fopen(fileName,"rb");
    if(file == NULL)
        return NULL;
    HamiltonianFile<T> Ha;
    fread(&Ha.signature,sizeof(char),2,file);
    fread(&Ha.ver_major,sizeof(unsigned short),1,file);
    fread(&Ha.ver_minor,sizeof(unsigned short),1,file);
    fread(&Ha.ver_build,sizeof(unsigned short),1,file);
    fread(&Ha.precision,sizeof(unsigned short),1,file);
    fread(&Ha.entry_type,sizeof(unsigned short),1,file);
    fread(&Ha.description,sizeof(char),1000,file);
    fread(&Ha.rowsNum,sizeof(int),1,file);
    fread(&Ha.colsNum,sizeof(int),1,file);
    fread(&Ha.elementsNum,sizeof(int),1,file);

    if(Ha.precision != sizeof(T)){
            printf("Data type of Hamiltonian matrix in the file is different with that in the program!\nexit.");
            exit(0);
    }

    CRS<T>* H=(CRS<T>*)malloc(sizeof(CRS<T>));
    H->rowsNum = Ha.rowsNum;
    H->colsNum = Ha.colsNum;
    H->elementsNum = Ha.elementsNum;

    H->rows = (int*)malloc(sizeof(int) * (H->rowsNum + 1));

    fread(H->rows,sizeof(int),H->rowsNum + 1,file);

    H->cols = (int*)malloc(sizeof(int) * H->colsNum);
    fread(H->cols,sizeof(int),H->colsNum,file);

    H->elements = (T*)malloc(sizeof(T) * H->elementsNum);
    fread(H->elements,sizeof(T),H->elementsNum,file);

    if(H->rows==NULL || H->cols==NULL || H->elements==NULL){
             printf("Failed to malloc memory!");
             exit(0);
    }

    return H;
}

template<class T>
CRS<T>* Hamiltonian::createCubeSingleElectron(HamiltonianInfo<T> &hi, MT::MersenneTwist& mt)
{
  int dim = 0;
  int XX = hi.getX();
  int YY = hi.getY();
  int ZZ = hi.getZ();
  T disorder = hi.getDisorder();
  if(XX > 1)
    dim += 1;
  if(YY > 1)
    dim += 1;
  if(ZZ > 1)
    dim += 1;

  int ROW_ELEMENTS=(dim * 2 + 1); //number of non-zero elements in each row.
  int H_ELEMENTS = (XX*YY*ZZ*ROW_ELEMENTS);

  CRS<T>* H = (CRS<T>*)malloc(sizeof(CRS<T>));
  H->elementsNum = H_ELEMENTS;
  H->rowsNum = XX*YY*ZZ;
  H->colsNum = H->elementsNum;
  H->elements = (T*)malloc(sizeof(T) * H->elementsNum);
  H->rows = (int*)malloc(sizeof(int) * (H->rowsNum + 1));
  H->cols = (int*)malloc(sizeof(int) * H->colsNum);

  int a[ROW_ELEMENTS];

  int x,y,z,i,j;
  for(z=0;z<ZZ;z++){
      for(y=0;y<YY;y++){
          for(x=0;x<XX;x++){
              a[0] = INDEX(x,y,z,XX,YY,ZZ);
              a[1] = INDEX(x-1,y,z,XX,YY,ZZ);
              a[2] = INDEX(x+1,y,z,XX,YY,ZZ);
              a[3] = INDEX(x,y-1,z,XX,YY,ZZ);
              a[4] = INDEX(x,y+1,z,XX,YY,ZZ);
              a[5] = INDEX(x,y,z-1,XX,YY,ZZ);
              a[6] = INDEX(x,y,z+1,XX,YY,ZZ);
              OrderData(a,ROW_ELEMENTS);

              H->rows[INDEX(x,y,z,XX,YY,ZZ)] = INDEX(x,y,z,XX,YY,ZZ) * ROW_ELEMENTS;
              for(i=0;i<ROW_ELEMENTS;i++){
                  if(a[i] == INDEX(x,y,z,XX,YY,ZZ)){
                      H->elements[INDEX(x,y,z,XX,YY,ZZ) * ROW_ELEMENTS + i] = on_site_potential<T>(mt,disorder);
                  }
                  else{
                      if(INDEX(x,y,z,XX,YY,ZZ) > a[i]){
                          T temp=0;
                          for(j = H->rows[a[i]];j<H->rows[a[i]+1];j++){
                              if(H->cols[j] == INDEX(x,y,z,XX,YY,ZZ)){
                                  temp = H->elements[j];
                              }
                          }
                          H->elements[INDEX(x,y,z,XX,YY,ZZ) * ROW_ELEMENTS + i] = temp;
                      }else{
                          H->elements[INDEX(x,y,z,XX,YY,ZZ) * ROW_ELEMENTS + i] = non_on_site_potential<T>(mt);
                      }
                  }
                  H->cols[INDEX(x,y,z,XX,YY,ZZ) * ROW_ELEMENTS + i] = a[i];
              }
           }
      }
  }

  H->rows[H->rowsNum] = H_ELEMENTS;

  normalize<T>(H, hi);
  return H;
}

template <typename T>
ELL<T>* Hamiltonian::readFromFile_ELL(char* fileName){
  CRS<T> c = Hamiltonian::readFromFile<T>(fileName);
  return Hamiltonian::convertToELLFromCRS<T>(c);
}

template <typename T>
void Hamiltonian::writeCRSToFile(CRS<T>* ha, char* prefix){
    char strElement[100]={0};
    char strCols[100]={0};
    char strRows[100]={0};

    sprintf(strElement,"%s_Elements",prefix);
    sprintf(strCols,"%s_Cols",prefix);
    sprintf(strRows,"%s_Rows",prefix);

    FILE* file = fopen(strElement,"w");
    fwrite(ha.elements,sizeof(T),ha->elementsNum,file);
    fclose(file);

    file = fopen(strCols,"wb");
    fwrite(ha.cols,sizeof(int), ha->colsNum, file);
    fclose(file);

    file = fopen(strRows,"wb");
    fwrite(ha.rows,sizeof(int), ha->rowsNum + 1, file);
    fclose(file);
}

template <class T>
CRS<T>* Hamiltonian::createBCCSingleElectron_CRS(HamiltonianInfo<T> &hi, MT::MersenneTwist& mt){
	int XX = hi.getX();
	int YY = hi.getY();
	int ZZ = hi.getZ();
	CRS<T>* ha = (CRS<T>*)malloc(sizeof(CRS<T>));

	const int colsNum = 9;

	ha->rowsNum = XX * YY * ZZ * 2;
	ha->colsNum = colsNum;
	ha->elementsNum = ha->rowsNum * ha->colsNum;

	ha->cols = (int*)malloc(sizeof(int) * ha->colsNum * ha->rowsNum);
	memset(ha->cols,0,sizeof(int) * ha->colsNum * ha->rowsNum);

	ha->elements = (T*)malloc(sizeof(T) * ha->colsNum * ha->rowsNum);
	memset(ha->elements,0,sizeof(T) * ha->colsNum * ha->rowsNum);

	ha->rows = (int*)malloc(sizeof(int) * (ha->rowsNum + 1));
	memset(ha->rows,0,sizeof(int) * (ha->rowsNum + 1));


	int x=0,y=0,z=0;
     for(z=0;z < ZZ * 2;z++){
    	 	 for(y=z%2; y < YY * 2; y+=2){
    	 		 for(x=z%2; x < XX * 2; x+=2){

    	 			 int ri = INDEX_BCC(x,y,z,XX,YY,ZZ);
    	 			 //each site has 8 neighbours
    	 			 int* c=&ha->cols[INDEX_BCC(x,y,z,XX,YY,ZZ) * colsNum];
    	 			 T* e=&ha->elements[INDEX_BCC(x,y,z,XX,YY,ZZ) * colsNum];

    	 			 int in=0;
    	 			 for(int i=-1;i<2;i+=2)
    	 				 for(int j=-1;j<2;j+=2)
    	 					 for(int k=-1;k<2;k+=2){
    	 						c[in] = INDEX_BCC(x+i,y+j,z+k,XX,YY,ZZ);
    	 						in++;
    	 					 }

    	 			 c[colsNum -1] = INDEX_BCC(x,y,z,XX,YY,ZZ);
    	 			 OrderData(c,colsNum);

    	 			 for(int i=0;i<colsNum;i++){
    	 				 if(c[i]==ri){
    	 					e[i] = on_site_potential<T>(mt,hi.getDisorder());
    	 				 }else{
    	 					e[i] = -1;
    	 				 }
    	 			 }

    	 			// printf("index:%d (%d,%d,%d)\n",INDEX_BCC(x,y,z,XX,YY,ZZ),x,y,z);
    	 		 }
    	 	 }
     }

     int temp=0;
     for(int i=0;i<ha->rowsNum+1;i++){
    	 	 ha->rows[i] = temp;
    	 	 temp += 9;
     }


     /*
     for(int i=0;i<ha->rowsNum;i++){
    	 	 for(int j=0;j<ha->colsNUm;j++){
    	 		 printf("%d\t",ha->cols[i * 9 + j]);
    	 	 }
    	 	 printf("\n");
    	 	 for(int j=0;j<ha->colsNUm;j++){
    	 		 printf("%.3f\t",ha->elements[i * 9 + j]);
    	 	 }
    	 	 printf("\n");
     }


     x=-1;y=-1,z=-1;
     printf("(%d,%d,%d)",x,y,z);
     printf("->(%d,%d,%d) index:%d\n",x,y,z, INDEX_BCC(x,y,z,XX,YY,ZZ));

     x=-1;y=-1,z=3;
     printf("(%d,%d,%d)",x,y,z);
     printf("->(%d,%d,%d) index:%d\n",x,y,z, INDEX_BCC(x,y,z,XX,YY,ZZ));

     x=-1;y=-1,z=1;
     printf("(%d,%d,%d)",x,y,z);
     printf("->(%d,%d,%d) index:%d\n",x,y,z, INDEX_BCC(x,y,z,XX,YY,ZZ));

     x=6;y=6,z=6;
     printf("(%d,%d,%d)",x,y,z);
     printf("->(%d,%d,%d) index:%d\n",x,y,z, INDEX_BCC(x,y,z,XX,YY,ZZ));

     x=4;y=4,z=6;
     printf("(%d,%d,%d)",x,y,z);
     printf("->(%d,%d,%d) index:%d\n",x,y,z, INDEX_BCC(x,y,z,XX,YY,ZZ));

     x=4;y=6,z=6;
     printf("(%d,%d,%d)",x,y,z);
     printf("->(%d,%d,%d) index:%d\n",x,y,z, INDEX_BCC(x,y,z,XX,YY,ZZ));
*/
     normalize<T>(ha, hi);
     return ha;
}


template<class T>
ELL<T>* Hamiltonian::createCubeSingleElectron_ELL(HamiltonianInfo<T> &hi, MT::MersenneTwist& mt)
{
  using namespace std;
  CRS<T>* c = Hamiltonian::createCubeSingleElectron<T>(hi, mt);
/*  std::cout<<"CRS format"<<std::endl;
  for(int i=0;i<20;i++){
  	std::cout<<c->elements[i]<<" ";
  }
  cout<<"------------------"<<endl;
  for(int i=0;i<20;i++){
  	std::cout<<c->cols[i]<<" ";
  }
  cout<<"------------------"<<endl;
  for(int i=0;i<20;i++){
  	std::cout<<c->rows[i]<<" ";
  }*/
  return Hamiltonian::convertToELLFromCRS<T>(c);
}

#endif
